Skip to content

[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294

Open
YWHyuk wants to merge 66 commits into
feature/togsim-cpp-tracefrom
fix/conv-wrapper-view-input
Open

[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294
YWHyuk wants to merge 66 commits into
feature/togsim-cpp-tracefrom
fix/conv-wrapper-view-input

Conversation

@YWHyuk

@YWHyuk YWHyuk commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

Fixes a conv codegen bug where the generated wrapper crashes with

padded_shape[3] += 2 * 0
IndexError: list index out of range

Root cause

A conv input node can be a ReinterpretView. call_kernel passes template args by
buffer name, so the wrapper is handed the base buffer, whose rank/shape may differ
from the view codegen assumed:

X
codegen (extract_info) X.layout.size = [1, 32, 8, 8] (the 4D view)
runtime (wrapper) X.shape = (1, 64, 32) (the base buffer)

The wrapper then inferred geometry from the runtime tensor (list(X.shape), X.shape[2],
X.shape[3]) and blew up.

It only fires when the view is free: a contiguous (B, N, C) buffer already is the
channels_last layout of (B, C, H, W), so inductor inserts no materializing copy. That is
exactly the transformer (B, N, C) -> (B, C, H, W) pattern — e.g. SegFormer's
efficient-attention spatial-reduction conv (Conv2d(C, C, k=sr, stride=sr), padding 0).
When the layout does not match, inductor inserts a copy kernel, a real 4D buffer arrives,
and the bug is invisible. That is why it went unnoticed.

Fix

Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the codegen
used into the wrapper and rebuild the logical NCHW inputs from it:

X = reinterpret_tensor(X, (1, 32, 8, 8), (0, 1, 256, 32), 0)
W = reinterpret_tensor(W, (16, 32, 2, 2), (128, 1, 64, 32), 0)
X_padding = torch.zeros((1, 32, 8, 8), device=X.device, dtype=X.dtype)
X_padding[:, :, 0:8, 0:8] = X

X.layout describes the data within the storage of the tensor that arrives (both come from
the same node), so the reinterpret is correct — and idempotent — whether the wrapper is
handed the base buffer, an already reinterpreted view, or a materialized copy.

The conv templates were the only place reading .shape at runtime; everything else already
derives geometry from the node layout at codegen time. No change to call_kernel,
mlir_argdefs, the MLIR signature, or arg_attributes
— the call site still passes the
raw base buffer and it now works.

Also fixed for free: the padding buffer was allocated via torch.zeros(shape) (always f32)
then .to(device=...) (CPU round-trip). It now allocates directly on device with X.dtype.

Verification

  • New regression test tests/ops/conv/test_conv_view_input.py: fails with the above
    IndexError on the pre-fix wrapper, passes after (max diff 2.3e-05 / 2e-04 vs CPU).
  • Normal conv path (real 4D buffer, padding=1) unchanged and numerically correct
    (max diff 3.8e-05 vs CPU); reinterpret_tensor is an identity there.

Follow-ups (not in this PR)

  • padding == 0 fast path: shapes are literals now, so the zeros alloc + full copy can be
    skipped for 1x1 / strided convs.
  • mlir_argdefs derives buffer_types numel from the base buffer, so a slicing view
    feeding e.g. gemm can still get a wrong MLIR arg size. Same class of bug, separate fix.
  • def_conv_kernel patches the signature with re.sub(r'(\d+)(?=xf32)', ...), hardcoding
    xf32; an f16 conv would not get its padded input size applied.

Found while checking #254 on this branch: SegFormer no longer hits the originally reported
Not supporting format error, but was blocked here instead.

YWHyuk and others added 30 commits July 7, 2026 16:09
… feed

Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime,
loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and
the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in
dma-fine-grained so multi-tile-K and conv loads do not collide; strip the
reduction accum marker from the memory_barrier slot.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into
togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs,
and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes;
the trace logs the first DRAM response and the assigned systolic array, and
scopes the compute barrier to its dispatch.
Default to the trace path; fix uninitialized Instruction fields, the matmul
accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and
dma_wait fusion; bound concurrent dispatches to the spad, round-robin
work-items within a partition, benchmark autotune and run the multi-tenant
scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose
[P3-trace] logs; fix the key.mlir compile race in load().
… runtime model

Replace the trace bridge's accumulated special cases with one dataflow rule and
clean up the runtime that consumes it.

Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all
current writers (occupancy=ISSUE when both are systolic-array ops, else
latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a
matmul that reads and writes the same buffer = a commutative accumulator): skip
its read edge and UNION its write, waiting only the non-matmul init seed and not
ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the
SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector
epilogue or the store wait every K matmul (fixes the pure-vector store that an
empty COMPUTE_BAR let slip).

Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA
drain), so the store JOINs the matmul writers directly. The whole emit/loader
chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the
runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a
stale producer fails loudly instead of emitting records the bridge would drop.
Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue).

Model compute-output spad footprint in the SRAM version/capacity machinery so
buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes
from the DMA records that touch the same buffer (a buf_bytes pre-pass); an
in-place buffer (accumulator, relu) is version-transparent so footprint is not
double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP
issue points (release before the COMP skip path so a skipped matmul still frees).

Runtime: collapse child_inst / _pipeline_children into one event-indexed
_deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot
release queue and the async-load wakeup into one _due_events timed-effect table
drained by process_due_events. Both are behavior-preserving (byte-identical).

Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at
init), and the round-robin disable mode is removed. Degenerate traces (a
consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather
than asserts that vanish under NDEBUG.

Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so
TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in
warns at command build.
… spad/2

The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner
carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace
is self-documenting. Both are read from togsim_runtime.h at codegen time -- the
version from the TOGSIM_ABI_VERSION define, the call-format text from a marked
block next to the declarations -- so they never drift when the ABI changes.
…rint

The bridge sums each work-item's distinct-buffer footprint (each buffer once, so
a reduction's reloads of the same section do not inflate it) onto its Tile.
can_issue then admits two concurrent dispatches only when each fits half the
spad; a dispatch larger than spad/2 runs alone with the whole spad, so two
work-items never compete for the shared spad and deadlock -- a runtime safety net
beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are
logged on the TILE_SCHEDULED trace line for debugging.
…election

select_tile fed the epilogue buffer count to gemm_combination_mapping as
max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers
were already freed and reusable by the epilogue. The codegen lays every buffer
out as a disjoint .spad global and never reuses freed space, so the estimate
undercounted the real footprint: for a fused matmul+relu kernel the budget came
to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088
B/lane including the relu output buffer. Under the full-spad guard this was
harmless (89088 < 131072), but the spad/2 guard rejected it and crashed
test_transformer_fusion with SpadOverflowError, since the heuristic path has no
tile-shrink fallback.

Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per
epilogue node plus one per extra read operand, matching what the codegen emits.
For the matmul+relu kernel the budget now equals the actual footprint exactly,
and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2.

Liveness-based spad reuse and broadcast over-allocation are tracked separately
as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane
scratchpad usage (spad_usage = stack_size + spad_size) and rejected the
tiling when that exceeded the spad/2 budget. But only the .spad section
actually lives in the scratchpad -- it is pinned there by the
--section-start=.spad link option. The kernel stack is in main memory
(sp is set up by pk in the -m region, not at the scratchpad vaddr), so it
does not consume the scratchpad and must not be charged against it. The
scalar frame is also shared across lanes, not per-lane, so adding it
double-counted.

On small configs (8x8) this falsely rejected feasible tilings: the
wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to
spare but were tripped over the spad/2 gate purely by the ~200-800 B
stack term, crashing compile with SpadOverflowError. Drop the stack
term; the .spad-only check still correctly rejects real buffer overflows
(e.g. sparsity, which is fixed separately by f05ac8a).
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP
padding entries are skipped instead of dereferencing uninitialized
garbage, fixing the host store segfault on 8-lane (wrapper3) configs
where a tile's split axis is not a multiple of vlane_stride*n_vu.
Add pytorchsim_functional_verify_per_kernel, a sub-option of
pytorchsim_functional_mode that localizes the first compiled kernel whose Spike
output diverges from a CPU reference.

When enabled, the generated wrapper compares every realized buffer (the output of
each fused kernel) against a CPU "golden" computed once per call by running the
original aten graph (V.graph.module) on CPU with the same inputs, via an
fx.Interpreter that records each node's output by name. A buffer is mapped to its
originating fx node through V.graph.get_buffer().origin_node, so the check reports
the kernel, the originating op, the offending indices and the max abs diff, then
raises at the first divergence (stop-at-first). Comparison granularity is the
fused-cluster output, the finest observable in a fused pipeline.

Auto-disabled when functional mode is off (no Spike values to verify); the config
accessor AND-gates the key with functional mode. Codegen bakes the
verify_init/verify_check calls into the wrapper only when enabled at compile time,
so clear the codegen cache when toggling. Tolerances via
TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4).

extension_functional_verify.py holds the graph registry and the runtime
golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level;
extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and
dropped the prologue's extra-read operands entirely, so a softmax-fused
attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The
codegen lays the softmax max/sum operands out as their own disjoint
weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so
the estimate undercounted the real footprint: on the 32x32 wrapper2 config
(16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted
kernel used 16896 B/lane and overflowed the budget by 512 B, crashing
test_transformer with SpadOverflowError. Under the 128x128 full-budget config
the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the
GEMM template (commit f05ac8a), now on the BMM prologue path.

Add an n_prologue_extra_read knob to gemm_combination_mapping that charges
each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N)
.spad buffer, matching what the codegen emits, and have the BMM template count
those operands the same way codegen does (the one numel-matching main input
reuses the matmul-operand buffer; every other read gets its own global). Tile
selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile
(8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose
check. The new parameter defaults to 0, leaving the GEMM and conv callers
unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism,
limitations, code map) and point to it from CLAUDE.md: a one-line entry in the
debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump
v1..v12 history was noise nobody reads.
[Frontend] Budget fused-prologue spad buffers in BMM tile selection
…s off

format_dma_inst_issued_trace_line and format_instruction_detail_line are only
ever used as the message argument to trace_instruction_line (a spdlog::trace
sink). Because the argument is evaluated eagerly at the call site, the kernel
builds a formatted std::string for every issued/finished instruction and then
spdlog drops it whenever the level is above trace (the default). Bail out of
both formatters when trace logging is disabled so the fmt::format work is
skipped entirely.

This removes wasted work but is not the simulation-speed bottleneck (that is the
per-cycle ready-queue rescan in Core::cycle, addressed separately); the conv
kernel wall time is unchanged within noise.
Core::cycle() walked every tile's ready-instruction list on every simulated
cycle to find one instruction to issue. Profiling an 8x8 conv kernel showed this
scan is ~96% of simulation time and, within it, the walk is >99%: the ready list
holds ~2000 instructions (mostly blocked on the SRAM-capacity / weight-slot
throttle), only one issues per cycle, and the blocked ones are never removed, so
the same ~2000 are re-examined every cycle (~2.6k list iterations/cycle, ~1e9
over 400k cycles). At small arrays this dominates because tiny tiles inflate both
the ready-list length and the number of DMA-wait stall cycles.

Gate the scan behind a per-Core _issue_dirty flag. The scan's outcome can only
change when the ready set grows or a resource frees, so set the flag on exactly
those events:
 - ready set grows: Instruction::dec_ready_counter, on enqueueing a now-ready
   instruction, sets the OWNING core's _issue_dirty via _owner_dirty_ref (wired
   in Core::issue when the tile is dispatched) -- per-core so a remote core's
   enqueue does not force every core to rescan;
 - SRAM freed (release_sram) and weight slot freed (apply_due) set _issue_dirty;
 - a new dispatch (issue) and a successful issue keep it set.
On a cycle with none of these the scan would re-walk the same blocked
instructions and issue nothing, so it is skipped. Instructions are never moved
between queues, so there is no re-admission churn.

Issue-identical: it never changes which instruction issues or when, only whether
the (side-effect-free under those conditions) scan runs. The 8x8 conv kernel
produces the same 403026 cycles as before, with its TOGSim wall dropping from
214.5s to 41.9s (~5x, 71% of scans skipped). A forced-scan invariant check (no
skipped cycle ever issues or inline-finishes a zero-cycle COMP) found 0
violations across 39 kernels spanning GEMM, BMM, softmax and conv. DMA-bound
kernels (the ones that hit the 6h CI cap) stall more and gain more.
In Core::cycle()'s issue scan, a COMP with compute_cycle == 0 finishes inline and
calls instructions.erase(it), but does not set `issued` or break -- so control
falls through to the `it++` at the end of the loop body, incrementing the
std::list iterator that erase() just invalidated. That reads a freed node and is
undefined behavior; it usually limps along because the freed node's next pointer
is briefly intact, but under a different allocator / sanitizer / build it can
crash or skip-or-duplicate the remaining ready instructions in that tile,
corrupting issue order.

Use the iterator erase() returns and continue, the standard erase-in-loop idiom.
No behavior change on the path that happened to work; it just removes the UB.
…X TOG

The main compile/sim path no longer generates or selects the legacy ONNX
Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv
(the build-skip now keys on trace.so), and TOGSimulator.run_standalone always
drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the
frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path
(extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR).

origins (which FX nodes a kernel came from) is preserved: logged per kernel run
and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy
ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at
the comment so the current parser is unaffected.

Also drop the dead tog_file param from mlir_gem5_compile_command, migrate
scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor
addresses and --attributes_list is no longer a Simulator option), and refresh
the CLAUDE.md TOG-generation notes.
Rework per-kernel codegen so the loop body is an ordered list of
load->compute->store steps (class Step + push_step + codegen_loops
iteration) instead of the fixed dma_loads/compute/dma_stores buffers.
A step may be compute-only or transfer-only; the compute_idx loop is
emitted only for steps that actually have compute content.

Drop the ad-hoc self.masks buffer: the reduction-tail mask (get_mask)
now emits into the compute buffer, since a mask is just compute.

Use the step model to express indirect (gather/scatter) access cleanly.
When the indirect index is produced in the same pass,
convert_indirect_indexing pushes a new step for the offset build, so the
index load, the offset build and the gather become separate steps bridged
through spad. This replaces the compute_dependecy -> dma_stores hack in
both convert_indirect_indexing and load(). push_step clones the CSE so the
name counter is shared but the dedup cache is per-step (each step is its
own compute_idx region). Validated: x[idx+2]+1 emits two steps and matches
CPU; add/matmul/reduce/softmax/layernorm and gather/scatter unchanged.
…ptor

Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.

The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).

Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
…ride, compute-loop multi-dim offset

Three refinements on top of the explicit offset descriptor:

1. Detect indirect access by an explicit symbol set instead of an "tmp"
   substring match. indirect_indexing now records str(index_var) in
   self.indirect_symbols; a _has_indirect(expr) helper tests the index
   free_symbols against it; the former "tmp"-string sites (store,
   get_dma_info, convert entry/indirect_dims/stride) use it. Removes the
   now-dead _find_indirect from lower_dma_to_gemmini.

2. Single indirect dim: pass the raw index spad and let the MVIN apply the
   gather stride per position (CONFIG4 offset_stride) instead of a
   vector_load/muli/vector_store round-trip that pre-multiplied the stride.
   emit_transfer carries offset_stride; decompose copies the attribute;
   lower programs CONFIG4 with it (was hardcoded to 1).

3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided
   indices, which a single CONFIG4 channel cannot do, so the sum stays in
   the kernel -- but build it in the compute loop (chunked by
   compute_vec_size, not a single tile-wide vector) and store it to a
   DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix)
   is not clobbered. push_step separates the offset-build loop from the
   gather that reads it.

Adds multi-dim gather and index-reuse cases to test_indirect_access.
Validated: indirect/scatter/embedding + the two new multi-dim cases pass;
add/matmul/softmax regression 0.
…emmini

memref.dma_start is a registered op with a fixed operand list, so it cannot carry the
runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape
values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate
the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers
directly to Gemmini.

- New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim
  collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's
  CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG
  rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather.
- togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx,
  dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without
  memref.dma_start's variadic tag indices.
- memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul
  lowering, erased by the merged lowering). No memref.* DMA ops remain.
- decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer
  (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its
  DRAM-grow half is what the future masked-DMA clamp replaces).
- dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer.

Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile +
async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv
(padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on
loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0
passes; wrong values, not a crash).
YWHyuk and others added 26 commits July 7, 2026 16:09
…dims, cleanup

Addresses the PR review of the masked-DMA non-dividing work:

- _masked_fill_bits detected a log-sum-exp reduction with `"exp" in origin.name`, which
  also matches expand / expm1 / experimental -- an ordinary non-dividing sum whose origin
  name contains "exp" would then get a -inf tail fill instead of the sum identity 0 and
  reduce to -inf. Match the op target instead (_origin_is_exp: the target's overload
  packet name == "exp"); same fix applied to the get_padding_type twin. Unit test in
  tests/lowering/test_masked_fill.py pins that expand/expm1 do not match.
- Remove the self._last_local_dims temporal coupling: get_dma_info now returns local_dims
  and load()/store() pass it to _masked_bounds explicitly, so a reordered or skipped
  get_dma_info can no longer feed stale clamp axes.
- Drop the underscore-prefixed locals in def_dma_op / store_epilogue (non-idiomatic),
  trim the oversized _masked_fill_bits / _masked_bounds docstrings to the invariant plus a
  link to docs/design/masked-dma-descriptor.md, and clarify the bf16 dtype-table comment
  (present only for table totality; bf16 is not actually supported).
Dropping the tile-divisibility recompile means non-dividing shapes rely entirely on the
masked-DMA clamp/fill; a gap is now a silent wrong answer, not a slow-but-safe recompile.
Add tests/ops/misc/test_masked_nondividing.py pinning that path: 2-D constant pad (per-dim
greedy pad recovery), sum over a broadcast operand (0 fill, guards the exp-substring bug),
softmax (genuine -inf fill), non-dividing elementwise, and non-dividing gather. Register it
in the CI allowlist (pytorchsim_test.yml). Also add the missing trailing newline to
test_indirect_access.py.
…ilogue

The ordered-steps refactor removed the reduction masking mechanism (self.masks /
get_mask / mask_cse) but left a compute_body.splice(self.masks) in
codegen_epilogue_body, so any matmul/bmm + reduction fusion crashed with
AttributeError: no attribute 'masks'. It was masked until now by the stale-spike
illegal-instruction CI failure. Remove the dead splice; test_matmul_reduction
(matmul+max, var_mean, residual+var_mean) passes. Non-dividing reduction-fusion
masking (padded lanes vs the reduction identity) stays a separate gap from that
refactor.
…ixel_shuffle trace)

Both the Gemmini descriptor and the TOGSim DMA model cap at 4 tile dims, but a
logical tile can exceed 4D (e.g. pixel_shuffle splits two spatial axes -> a 5D
tile [1,1,2,4,2]). lower_transfer_to_gemmini reduces >4D as part of emitting
Gemmini asm, but that runs only on the Spike/gem5 lowering path. The trace
producer (build_tog) reads togsim.transfer BEFORE that, so it emitted a 5D DMA
and TOGSim rejected it: "issued tile is not supported format.. tile.size: 5".
Exposed by the pixel_shuffle case newly added to the wrapper1 CI suite.

Add a peel_transfer POST_OPT pass that reduces every >4D togsim.transfer up front
(before build_tog) while KEEPING it a togsim.transfer, so both consumers see <=4D.
It mirrors lower_transfer_to_gemmini's two reductions: collapse unit dims when the
effective rank is <=4 (memref.collapse_shape), else peel the outer effective dims
into an affine.for nest with the lane-banked physical SRAM offset. The gemmini
lowering keeps its own reduction as a now-defensive no-op.

Validated on 128x128: test_floormod_axis_split (pixel_shuffle + 9 others),
test_masked_nondividing, test_matmul, test_conv2d all pass.
…-safe)

_GoldenInterpreter recorded each node output on CPU but returned the original
(npu) tensor to downstream nodes, and cast the recorded value to float32. A
device-baked op (e.g. arange on npu) or a consumer that requires operand device
agreement (aten.index needs its index tensors on the indexed tensor's device)
then saw a mixed CPU/npu pair and the whole golden threw, silently disabling the
per-kernel verify for that graph -- e.g. MobileNet's depthwise conv lowered to
constant_pad + index gather.

run_node now moves EVERY tensor output to CPU (preserving dtype) and returns it,
so downstream nodes see CPU operands; only the recorded golden is cast to float32
for the allclose compare, so int index tensors stay usable.
The cat template's _calculate_input_tile_sizes handed each input a
concat-dim tile of min(extent, remaining_spad_budget). When that budget
did not cover an input's full concat extent, the tile width could fail
to divide the extent (e.g. extent 64, tile 62 on the 8-lane config).
The copy loop steps by that tile over [0, extent), so a non-dividing
tile produces a ragged final iteration with no remainder mask: the MVIN
over-reads past the input and the MVOUT over-writes past the output,
corrupting the result.

This surfaced on the 8x8 (wrapper3) CI config, where the per-input spad
budget lands at 126 for two 64-wide inputs, truncating the second tile
to 62. It scrambled ~45% of the cat output and broke LlamaDecoderLayer
(rotary rotate_half cat) and the diffusion UNet skip-connection concat.
Larger arrays give a bigger budget so both inputs got the full 64-wide
(dividing) tile, which is why 32x32 and 128x128 passed.

Snap each input's concat-dim tile down to the largest divisor of its
extent that fits the budget, so every loop iteration is full-width and
in-bounds. When the budget already covers the full extent the tile is
unchanged, so 32x32/128x128 codegen is a no-op.

Verified with a minimal cat repro (64+64 -> 128 on 8x8: fails before,
diff 0 after) and the real test_llama LlamaDecoderLayer at 8x8 (now
passes).
The vlane index_expr lowering reconstructs each lane element's logical
index along the split axis as stride_dim + outer_dim * nr_vector_lane,
where outer_dim is the sublane-row index. The next sublane row starts
vector_lane * vlane_stride elements later (vector_lane lanes, each
holding vlane_stride positions), but the multiplier used nr_vector_lane
alone, dropping the vlane_stride factor. Any index tensor (e.g.
arange/iota) whose split axis spills into a second sublane row then had
every row past the first come out shifted, producing wrong values.

This fired on the 8x8 (wrapper3) config: arange(32) at 8 lanes with
vlane_stride 2 needs two sublane rows, so positions 16-31 came out as
p-8. Larger arrays keep the 32-wide axis in a single row (outer_dim
always 0), which is why 32x32/128x128 were unaffected. It broke
test_llama's LlamaModel (rotary position ids) at 8x8.

Multiply the outer/row term by vector_lane * vlane_stride. This equals
the old value whenever vlane_stride == 1, so single-row splits and
larger configs are a codegen no-op.

Verified with a minimal arange(32) repro at 8x8 (positions 16-31 wrong
before, diff 0 after) and the full test_llama LlamaModel at 8x8 (now
passes).
init_tile_size gave a 1-D tile of 2*vlane_stride*vector_lane regardless
of the data extent. When the extent is smaller than that tile the
reduction/pointwise MVIN reads past the end of the DRAM buffer, and the
compute loop folds the out-of-range elements into the result. For a sum
reduction that means adjacent-DRAM garbage is squared/added into the
sum.

This broke test_single_perceptron's Loss on the 128x128 config: the MSE
mean reduces 128 elements but the tile was 2*2*128 = 512, so the MVIN
over-read 384 elements past the 128-wide input and roughly doubled the
squared-error sum (loss diff ~4593). It fired on wide-lane configs
(tile 512 > extent 128); 32x32 gives tile 128 == extent and 8x8 gives a
tile that divides 128, so both were unaffected.

Cap the 1-D tile to min(2*vlane_stride*vector_lane, extent) so the tile
never exceeds the data. It is a no-op whenever extent >= tile, so larger
reductions and other configs are unchanged.

Verified with a minimal mse_loss(1-D) repro at 128x128 (over-reads
before -> diff 0 after, across extents 64/100/127/200/300) and the real
test_single_perceptron at 128x128 (Loss now passes).
torch.sort's frontend lowering exposes the sorted VALUES as the sole
scheduler-visible template output; the INDICES buffer is written in place
by the same kernel (make_inplace) but is not registered as an output. For
argsort the values output is dead, so Inductor DCEs the whole sort kernel
and the indices buffer is returned uninitialized (all-zeros), silently
corrupting any argsort result (e.g. DeepSeek MoE expert routing).

Register the in-place indices write as a MutationOutput owned by the sort
operation so the scheduler keeps the kernel alive whenever indices is
consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small
MLIRTemplateBuffer subclass that can own extra MutationOutputs and have
MLIRTemplateCaller.output_node build it; the sort lowering attaches the
indices mutation via add_mutation_output. For templates with no mutation
this is behaviorally identical to the base (get_outputs == [self]).

Then gate the values MVOUT on the values output being live: in argsort it
is pruned from the kernel args, so emitting its MVOUT referenced an
undeclared %YV. The values scratchpad is still produced internally for the
bitonic compare; only the dead DRAM store is skipped.

Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort
with both outputs live is unchanged; test_sort, test_matmul, test_cat,
test_indirect_access unaffected by the output_node change. Note: x[argsort]
end to end additionally needs the indirect-gather codegen fix (separate).
convert_indirect_indexing moves an indirect index into an offset spad and
zeros the indirect term in the DRAM base-offset expression, but it only
walked index.args. A BARE indirect index (x[idx], where the load index IS
the symbol so index.args is empty) was left unzeroed, so the DMA base
offset's affine.apply referenced the loop-local index SSA value before it
was defined ("use of undeclared SSA value name"). Any 1-D gather with a
runtime index tensor failed to compile at codegen; it was masked for
argsort only because the sort DCE bug left the indices all-zero (a
degenerate constant load).

Subs any remaining indirect symbol to 0 after the arg loop (mirrors the
store path). The per-position gather is carried by the offset spad, so the
base offset is correctly 0.

Verified: x[idx] (N=32/64) and x[x.argsort()] now compile and match CPU;
test_indirect_access.py (incl. scatter/index_add and multi-dim) and
test_sort.py unchanged.
_masked_bounds reverse-engineered per-dim padding from the single flat
index offset and clamped each padded load to [pad, pad + input_extent).
That recovery is ill-posed: the offset mixes the tap shift with the
padding, and under channels_last the stride-1 channel axis absorbs the
offset remainder, producing an impossible clamp (e.g. [4, 8) on a size-2
channel tile). Spike then zeroed the whole load, making channels_last
depthwise conv 99.9 percent wrong while NCHW was fine.

The clamp was also unnecessary: at pad positions the consumer already
yields the correct value (a compute-side select for a padded gather, or
the pad op's own fill), so reading OOB is harmless. Keep only the ragged
tail clamp (glo=0, ghi=loop_extent), which is genuinely load-bearing for
non-dividing tiles.

Verified on 32x32: channels_last depthwise conv now matches CPU; NCHW
conv and standalone ragged F.pad still pass. Broader regressions (padded
reduction loads, full MobileNet) to be addressed separately.
codegen_epilogue_body() decided whether to emit the template buffer's MVOUT
by asking "did a fused epilogue already store something?":

    if len(self.stores._lines) == 0:
        template_store()

That is a proxy for the real question, "does anything outside this kernel still
read the template buffer?", and it is wrong in two of the four cases:

  - epilogue + live template buffer  -> store skipped, buffer never written
  - no epilogue + dead template buffer -> store emitted for a pruned DRAM arg

The first case corrupts DeepSeek-V3's MoE gate. The gate is
sigmoid(mm(hidden, gate_w) + bias), so add+sigmoid fuse onto the GEMM template,
but the raw mm result is also read by a later gather kernel. Its buffer was
declared as a kernel output and never stored, so the gather read an
uninitialized buffer (all zeros), scrambling expert routing. Every config was
affected; only 8x8 crossed the test's loose atol=2e-1 (final Max abs diff
0.2546) while 32x32 landed just under it and passed.

Ask the real question instead. Inductor already answers it:
Kernel.remove_kernel_local_buffers() drops a buffer only when
Scheduler.can_buffer_be_removed_through_fusion() finds every user inside this
fused kernel. A buffer that survived removal therefore has an outside user and
must be stored. Record the template buffer's name in def_kernel() and
def_conv_kernel() (conv has its own implementation) and gate template_store()
on it, in both the main and the reduction_fusion branches. The proxy is gone.

Use the kernel-local self.removed_buffers, not V.graph.removed_buffers: the
former is what remove_buffer() populates, and the latter is still empty when
the store hook runs.

Note this makes a previously missing MVOUT issue, so DRAM traffic and cycle
counts rise for any model with a live template buffer under a fused epilogue.
That is the correct cost, not a regression, but reported cycles will shift.

Verified on the 8x8 config: test_deepseek_v3_base goes from Max abs diff 0.2546
to passing; a dead template buffer (relu(mm)) still prunes its DRAM arg and
emits a single MVOUT; the trace/TOGSim timing path runs with the extra store.
Regressions pass: test_bmm_reduction, test_matmul, test_bmm, test_conv2d,
test_cnn, test_group_conv, test_pool, test_reduce, test_softmax, test_layernorm,
test_mlp.
Implement abs, sign, isnan, isinf, fmod, bitwise shifts, copysign, erfc,
hypot, cosh, sinh, acos, asin, atan, atan2, acosh, asinh, atanh, and log/
atan (VCIX) in the Inductor MLIR backend. Float-only transcendentals cast
non-float inputs to f32 while leaving native float widths (f16/f64)
untouched.
Port the C++ MathLog/MathAtanToVCIX patterns to the in-process Python
VCIX pass: add math.log (opcode 2, imm 2) and math.atan (opcode 2, imm 3)
to _MATH_VIV and MARKERS. Encoding matches the spike rs1 / gem5 VS1
sub-opcodes (sin=0, cos=1, log=2, atan=3). No C++ mlir-opt pass needed.
Add the pointwise op test under tests/ops/elementwise and register it in
the test workflow allowlist.
Bump the third-party pins to the releases shipping the new pointwise
instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
reduction_init returned "-inf"/"inf" for max/min regardless of dtype, so an
integer max/min reduction fed an invalid inf identity into every consumer:
the accumulator init, the template reduction init, and the masked-DMA tail
fill (torch.tensor(inf, dtype=int64) overflowed -- surfaced compiling swinv2
with batch>1). reduction_combine_vec also hardcoded the float multi_reduction
kinds <maximumf>/<minimumf>, which reject integer vectors.

Return the dtype's representable extreme from reduction_init for integer
max/min (max identity -inf -> iinfo.min, min identity +inf -> iinfo.max;
bool -> 0/1), and emit the signed-integer multi_reduction kinds <maxsi>/<minsi>
for integer element types. sum/prod need no change: the <add>/<mul> combining
kinds are polymorphic and ops.add/ops.mul already select arith.add{i,f} by
dtype (as ops.maximum/minimum already select maxsi vs maximumf).

Verified: int amax/amin/sum/prod reductions compile and match CPU; float
test_reduce (ReduceMax) and test_softmax unchanged. Also clears the swinv2
batch>1 masked-fill overflow (it then hits a separate unsupported
shifted-window view, tracked separately).
…lar DMA index

torch.roll has no Inductor lowering; it decomposes to index_select(fmod(...)) -- a
cyclic-shift gather whose ModularIndexing DMA index the affine-only descriptor
cannot express ("Unlinearized floor/mod in DMA index"). Rewrite it as the two
contiguous halves the shift produces, stitched by cat:

    roll(x, s, d) == cat([slice(x, N-s, N), slice(x, 0, N-s)], d)

each half is a plain strided slice with no modulo. Do it as a lowering (removing
roll from the decomp table so it is not decomposed first) rather than a
decomposition, so we can also REALIZE the input with copy_input: roll's producer
is often a reshaped view (e.g. swinv2's window_reverse), and the slice offset would
otherwise fuse into that reshape's modular index, re-creating a shifted (p+c)%m
form. A plain .contiguous()/clone does NOT create the buffer boundary (Inductor
inlines it); copy_input does.

Verified: torch.roll (1-D / multi-dim / arbitrary shift) and roll+reduction match
CPU; swinv2 batch>1 now compiles through the shifted-window attention (previously
blocked by the roll's cyclic shift and then by the fused slice offset).
A composition of aligned reshapes on one iteration variable (e.g. swinv2's window
partition fused with a later reshape) leaves a nested index like
ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor
simplify_with_ranges reduces. collect_boundaries then skips its cut points (the
inner base is not a bare variable) and the affine-only DMA check rejects it.

Add a general, pattern-free canonicalizer (flatten_nested_floormod / _as_digit):
any nesting of FloorDiv/ModularIndexing over a single symbol is a digit extractor
(v // A) % M and collapses to one level by composing divisors, from four
divisibility-guarded algebraic identities. Applied in collect_boundaries and in
_fold_with_ranges. Multi-variable inner arguments (e.g. a roll shift v+c) are left
untouched.

Verified: numeric equivalence on the swinv2 window index; test_reduce, test_layernorm,
test_cat, test_indirect_access unaffected.
The wrapper header imported empty_strided but not empty_strided_cpu, so a graph
that allocates a CPU-side buffer (e.g. swinv2's shifted-window attention mask,
which Inductor keeps on CPU) failed at runtime with
"NameError: name 'empty_strided_cpu' is not defined". Add the same binding the
default Inductor PythonWrapperCodegen header provides.
…yout

A per-lane reduction is lowered as a 2-D [reduction | batch] collapse:
multi_reduction reshapes the accumulator to <reduction_numel x red_size>
and reduces axis 0, which assumes the reduction axis is the outermost
in-lane run and every non-reduction (batch) axis is inner. get_dma_info
reorders the tile axes for reduction loads to guarantee that, but only
for the 2-D and 3-D cases. The 4-D case mis-tested is_reduction
(reduction_depth < 3, false once there are three batch dims) and then
raised NotImplementedError, and rank 5+ fell through to the default
row-major order -- both leaving the reduction axis innermost, so the
reduce collapsed a batch axis instead of the reduction axis.

This is hit whenever a non-reduction dim-merge is blocked and an extra
batch axis stays in-lane: SwinV2 cosine window attention adds a gathered
relative-position bias (table[idx[q,k], head]) before the softmax amax,
keeping head separate from query (4-D tile, head bled into the max), and
its post-attention LayerNorm var_mean splits the token axis into several
loop vars (6-D tile, mean off by ~0.005).

Generalize the reorder to any tile rank: a tile that carries the
reduction axis places that axis-group outermost via
range(r, L) + range(r-1, -1, -1), which reduces to the existing 3-D
order for L=3. Adds test_reduce_gather_bias (fails max abs diff ~3.3
before the fix).
SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA
index" because the shifted-window path (torch.roll composed with window
partition/reverse) produced views axis-split could not linearize. With
the roll->slice+cat lowering, the nested/shifted mod handling, and the
reduction-axis layout fix, the whole backbone now compiles and matches
CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2).

Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to
the other model tests, and list SwinV2 in the README model coverage.
conv_multi_tile_mapping (used for batch > 1 convs) collects every tile
that fits SPAD into tile_candidates and returns them sorted, but it
guarded the "no mapping" error on max_used_spad_size instead of on
tile_candidates. max_used_spad_size is only bumped when a candidate also
satisfies max_k_h_w <= k_h, and max_k_h_w is initialized to K_W, so it
only ever fires for the full-kernel (k_h == K_H) tile. When that tile
overflows SPAD -- e.g. a CLIP ViT-B/32 patch conv (3->768, 32x32 stride
32) at batch > 1 -- the guard raised "Cannot find a valid mapping" even
though smaller-k_h tiles fit and were already in tile_candidates. The
mapping variable it tracks is never returned.

Guard on `not tile_candidates` instead, so the conv is rejected only
when no tile fits at all. The returned candidate list is unchanged, so
convs that already worked are unaffected. Adds a batched patch-conv case
to test_conv2d.py. Fixes issue #252.
With the conv-mapping fix, CLIP's vision transformer runs with batch > 1
and matches CPU end to end (max diff ~1e-5). Add tests/models/test_clip.py
(CLIPVisionModel, batch 2, patch 32), wire it as a self-hosted CI job,
and list CLIP in the README model coverage.
A conv input node can be a ReinterpretView. call_kernel passes template args
by buffer *name*, so the wrapper receives the base buffer, whose rank/shape may
differ from the view codegen assumed. The wrapper then did

    padded_shape = list(X.shape)
    padded_shape[3] += 2 * PADDING_W        # IndexError on a 3D base buffer

This fires whenever the view is free: a contiguous (B, N, C) buffer already is
the channels_last layout of (B, C, H, W), so inductor inserts no materializing
copy. That is exactly the transformer (B, N, C) -> (B, C, H, W) pattern, e.g.
SegFormer efficient attention.

Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the
codegen used into the wrapper and rebuild the logical NCHW inputs from it, so the
wrapper is correct whether it is handed the base buffer, an already reinterpreted
view, or a materialized copy (reinterpret is idempotent in all three cases). The
conv templates were the only place reading .shape at runtime; everything else
already derives geometry from the node layout at codegen time.

No change to call_kernel, mlir_argdefs, the MLIR signature, or arg_attributes.

Also allocate the padding buffer directly on the input device with the input
dtype: torch.zeros() defaulted to f32 and to(device=...) round-tripped via CPU.
Covers a conv fed by a free ReinterpretView of a lower-rank buffer. Fails with
IndexError on the pre-fix wrapper and passes after it.
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch 4 times, most recently from d9131ba to 1ff2ebb Compare July 10, 2026 13:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants